package com.example.search.utils;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.example.search.model.Page;
import org.apache.commons.lang.StringUtils;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.get.MultiGetRequest;
import org.elasticsearch.action.get.MultiGetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.MultiSearchRequest;
import org.elasticsearch.action.search.MultiSearchResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.*;

import static org.elasticsearch.index.query.QueryBuilders.rangeQuery;
import static org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders.weightFactorFunction;

@Component
public class EsRestClientUtil {

    private static RestHighLevelClient client;

    @Autowired
    public EsRestClientUtil(RestHighLevelClient client) {
        this.client = client;
    }


    //==================== 保存 ↓ ==============================

    public static IndexResponse save(String index, String type, String id, Map<String, Object> source) throws IOException {
        IndexRequest indexRequest = new IndexRequest(index, type, id);
        indexRequest.source(source);
        return client.index(indexRequest, RequestOptions.DEFAULT);
    }

    public static IndexResponse save(String index, String type, Map<String, Object> source) throws IOException {
        String id = null;
        if (source.containsKey("id")) {
            Object idObj = source.get("id");
            id = idObj == null ? UUID.randomUUID().toString().replace("-", "") : String.valueOf(idObj);
        }
        IndexRequest indexRequest = new IndexRequest(index, type, id);
        indexRequest.source(source);
        return client.index(indexRequest, RequestOptions.DEFAULT);
    }

    public static int saveBatch(String index, String type, List<Map<String, Object>> sourceList) throws IOException {
        int count = 0;
        if (!CollectionUtils.isEmpty(sourceList)) {
            String id = null;
            for (Map<String, Object> source : sourceList) {
                id = UUID.randomUUID().toString().replace("-", "");
                IndexRequest indexRequest = new IndexRequest(index, type, id);
                indexRequest.source(source);
                client.index(indexRequest, RequestOptions.DEFAULT);
            }
        }
        return count;
    }


    //==================== 更新 ↓ ==============================

    public static UpdateResponse update(String index, String type, String id, Map<String, Object> source) throws IOException {
        UpdateRequest updateRequest = new UpdateRequest(index, type, id);
        updateRequest.doc(source);
        return client.update(updateRequest, RequestOptions.DEFAULT);
    }

    public static Object saveOrUpdate(String index, String type, String id, Map<String, Object> source) throws IOException {
        Object result = null;
        if (exists(index, type, id)) {
            result = update(index, type, id, source);
        } else {
            result = save(index, type, id, source);
        }
        return result;
    }


    //==================== 判断是否存在 ↓ ==============================

    public static boolean exists(String index) throws IOException {
        return exists(index, null);
    }

    public static boolean exists(String index, String type) throws IOException {
        return exists(index, type, null);
    }

    public static boolean exists(String index, String type, String id) throws IOException {
        GetRequest getRequest = new GetRequest(index);
        if (StringUtils.isNotBlank(type)) {
            getRequest.type(type);
        }
        if (StringUtils.isNotBlank(id)) {
            getRequest.id(id);
        }
        return client.exists(getRequest, RequestOptions.DEFAULT);
    }


    //==================== 数量 ↓ ==============================

    /*public static void count(String index, String type) throws IOException {
        //CountRequest countRequest = new CountRequest(index, type);
    }*/


    //==================== 查询 ↓ ==============================

    public static GetResponse queryById(String index, String type, String id) throws IOException {
        GetRequest getRequest = new GetRequest(index, type, id);
        return client.get(getRequest, RequestOptions.DEFAULT);
    }

    private static SearchResponse queryPage(String[] index, String[] type, Page page) throws IOException {
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.from(page.getStartRow());
        sourceBuilder.size(page.getPageSize());
        if (!CollectionUtils.isEmpty(page.getTermQueryMap())) {
            Map<String, Object> conditionMap = page.getTermQueryMap();
            for (Map.Entry<String, Object> entry : conditionMap.entrySet()) {
                sourceBuilder.query(QueryBuilders.termQuery(entry.getKey(), entry.getValue()));
            }
        }
        if (!CollectionUtils.isEmpty(page.getSortMap())) {
            Map<String, SortOrder> sortMap = page.getSortMap();
            for (Map.Entry<String, SortOrder> entry : sortMap.entrySet()) {
                sourceBuilder.sort(entry.getKey(), entry.getValue());
            }
        }

        SearchRequest request = new SearchRequest(index);
        if (type != null && type.length > 0) {
            request.types(type);
        }
        request.source(sourceBuilder);
        return client.search(request, RequestOptions.DEFAULT);
    }

    public static SearchResponse queryPageByIndexAndType(String[] index, String[] type, Page page) throws IOException {
        return queryPage(index, type, page);
    }

    public static SearchResponse queryPageByIndex(String[] index, Page page) throws IOException {
        return queryPageByIndexAndType(index, null, page);
    }

    public static List<JSONObject> queryJsonByIndexAndType(String[] index, String[] type, Page page) throws IOException {
        SearchResponse response = queryPageByIndexAndType(index, type, page);
        SearchHits hits = response.getHits();

        List<JSONObject> list = new ArrayList<>();
        if (hits != null) {
            SearchHit[] hitsArr = response.getHits().getHits();
            for (SearchHit searchHit : hitsArr) {
                list.add(JSON.parseObject(JSON.toJSONString(searchHit), JSONObject.class));
            }
        }
        return list;
    }

    public static List<JSONObject> queryJsonByIndex(String[] index, Page page) throws IOException {
        return queryJsonByIndexAndType(index, null, page);
    }

    public static SearchResponse queryPageByRange(String[] index, String[] type, Page page) throws IOException {
        Map<String, Object[]> rangeMap = page.getRangeQueryMap();
        Map<String, Object> termQueryMap = page.getTermQueryMap();
        FunctionScoreQueryBuilder.FilterFunctionBuilder[] rangeBuilder = null;
        int count = 0;

        if (!CollectionUtils.isEmpty(rangeMap)) {
            rangeBuilder = new FunctionScoreQueryBuilder.FilterFunctionBuilder[rangeMap.size()];
            if (!CollectionUtils.isEmpty(rangeMap) && !CollectionUtils.isEmpty(termQueryMap)) {
                rangeBuilder = new FunctionScoreQueryBuilder.FilterFunctionBuilder[rangeMap.size() + termQueryMap.size()];
            }
            String rangeName = null;
            Object rangeValLeft = null;
            Object rangeValRight = null;

            for (Map.Entry<String, Object[]> entry : rangeMap.entrySet()) {
                rangeName = entry.getKey();
                rangeValLeft = entry.getValue()[0];
                rangeValRight = entry.getValue()[1];

                rangeBuilder[count] = new FunctionScoreQueryBuilder.FilterFunctionBuilder(
                        rangeQuery(rangeName).from(rangeValLeft).to(rangeValRight), weightFactorFunction((count + 1))
                );
                count++;
            }
        }

        if (!CollectionUtils.isEmpty(termQueryMap)) {
            String sourceName = null;
            Object sourceVal = null;
            if (count < 1) {
                rangeBuilder = new FunctionScoreQueryBuilder.FilterFunctionBuilder[termQueryMap.size()];
            }
            for (Map.Entry<String, Object> entry : termQueryMap.entrySet()) {
                sourceName = entry.getKey();
                sourceVal = entry.getValue();

                rangeBuilder[count] = new FunctionScoreQueryBuilder.FilterFunctionBuilder(
                        QueryBuilders.termQuery(sourceName, sourceVal), weightFactorFunction((count + 1))
                );
                count++;
            }
        }

        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.from(page.getStartRow());
        sourceBuilder.size(page.getPageSize());
        if (null != rangeBuilder) {
            sourceBuilder.query(QueryBuilders.functionScoreQuery(rangeBuilder));
        }
        SearchRequest searchRequest = new SearchRequest(index);
        if (null != type && type.length > 0) {
            searchRequest.types(type);
        }
        searchRequest.source(sourceBuilder);
        return client.search(searchRequest, RequestOptions.DEFAULT);
    }

    public static SearchResponse queryPageByRange(String[] index, Page page) throws Exception {
        return queryPageByRange(index, null, page);
    }

    public static SearchResponse queryByBuilder(String[] index, SearchSourceBuilder sourceBuilder) throws Exception {
        SearchRequest searchRequest = new SearchRequest(index);
        if (null != sourceBuilder) {
            searchRequest.source(sourceBuilder);
        }
        return client.search(searchRequest, RequestOptions.DEFAULT);
    }

    public static MultiGetResponse mget(List<MultiGetRequest.Item> requests) throws Exception {
        MultiGetRequest request = new MultiGetRequest();
        if (!CollectionUtils.isEmpty(requests)) {
            requests.stream().forEach(request::add);
        }
        return client.mget(request, RequestOptions.DEFAULT);
    }

    public static MultiSearchResponse msearch(List<SearchRequest> requests) throws Exception {
        MultiSearchRequest request = new MultiSearchRequest();
        if (!CollectionUtils.isEmpty(requests)) {
            requests.stream().forEach(request::add);
        }
        return client.msearch(request, RequestOptions.DEFAULT);
    }
}
